package it.tukano.swingdeco.tree;

import it.tukano.fun.Function1;
import it.tukano.fun.Function2;
import it.tukano.fun.OptionalValue;
import it.tukano.swingdeco.containers.JScrollPaneExt;
import it.tukano.swingdeco.event.InputAdapter;
import it.tukano.swingdeco.event.MouseAdapterExt;
import it.tukano.swingdeco.event.MouseEventExt;
import java.awt.Component;
import java.awt.Point;
import java.awt.Rectangle;
import java.awt.event.MouseListener;
import java.util.Enumeration;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.List;
import javax.swing.JLabel;
import javax.swing.JTree;
import javax.swing.event.TreeSelectionEvent;
import javax.swing.event.TreeSelectionListener;
import javax.swing.tree.DefaultMutableTreeNode;
import javax.swing.tree.DefaultTreeCellRenderer;
import javax.swing.tree.DefaultTreeModel;
import javax.swing.tree.DefaultTreeSelectionModel;
import javax.swing.tree.TreeModel;
import javax.swing.tree.TreeNode;
import javax.swing.tree.TreePath;

/**
 * A jtree extension
 * @author pgi
 */
public class JTreeExt extends JTree implements Iterable<DefaultMutableTreeNode> {

    private static final long serialVersionUID = 0;

    private static DefaultTreeModel newModel() {
        DefaultMutableTreeNode root = new DefaultMutableTreeNode("root");
        DefaultTreeModel model = new DefaultTreeModel(root);
        return model;
    }

    private MouseListener selectionCleaner = new MouseAdapterExt() {

        @Override
        protected void leftMousePressedExt(MouseEventExt ext) {
            TreePath path = getPathForLocation(ext.getX(), ext.getY());
            if(path == null) {
                setSelectionPath(null);
            }
        }
    };
    private JScrollPaneExt scroller;

    /* Listeners added to this tree via addTreeSelectionListener and related */
    private final List<TreeSelectionListener> LISTENERS = new LinkedList<TreeSelectionListener>();

    /**
     * Instance initializer
     */
    public JTreeExt() {
        super(newModel());
        setRootVisible(false);
    }
    
    public JTreeExt onDoubleClick(final Function1<DefaultMutableTreeNode, ?> callback) {
        addMouseListener(new InputAdapter() {

            @Override
            protected void onLeftMousePressed(MouseEventExt e) {
                if(e.isDoubleClick()) {
                    OptionalValue<DefaultMutableTreeNode> node = getNodeAt(e.getPoint());
                    if(node.isDefined()) {
                        callback.apply(node.get());
                    }
                }
            }
        });
        return this;
    }

    public JTreeExt setOneTouchExpandable() {
        setToggleClickCount(1);
        return this;
    }

    /**
     * Returns a JScrollPane wrapping this tree
     * @return a JScrollPane fro this tree
     */
    public JScrollPaneExt getScroller() {
        return scroller == null ? (scroller = new JScrollPaneExt(this)) :  scroller;
    }

    /**
     * Set the DefaultTreeModel of this JTreeExt.
     * @param newModel a DefaultTreeModel instance. Throws an exception if
     * newModel is not of DefaultTreeModel type
     */
    @Override
    public void setModel(TreeModel newModel) {
        if(newModel instanceof DefaultTreeModel) {
            super.setModel(newModel);
        } else {
            throw new UnsupportedOperationException("model should be a DefaultTreeModel");
        }
    }

    /**
     * Returns the root of this tree
     * @return the root of this tree
     */
    public DefaultMutableTreeNode getRoot() {
        return (DefaultMutableTreeNode) getModel().getRoot();
    }

    /**
     * Add a node with the given object to the tree root
     * @param userObject the user object to add
     * @return this
     */
    public JTreeExt append(Object userObject) {
        DefaultTreeModel model = getTreeModel();
        DefaultMutableTreeNode root = getRoot();
        DefaultMutableTreeNode node = new DefaultMutableTreeNode(userObject);
        model.insertNodeInto(node, root, root.getChildCount());
        return this;
    }

    /**
     * Add a listener that will be notified when the selection changes
     * @param callback the callback
     * @return this
     */
    public JTreeExt onSelectionChange(final Function1<DefaultMutableTreeNode, ?> callback) {
        addTreeSelectionListener(new TreeSelectionListener() {

            public void valueChanged(TreeSelectionEvent e) {
                callback.apply(getSelectedNode().get());
            }
        });
        return this;
    }

    /**
     * Set the tree to clear selection when the user pressed the mouse in a
     * region of the tree that doesn't contain items
     * @return this
     */
    public JTreeExt setClearSelectionOnClick() {
        addMouseListener(selectionCleaner);
        return this;
    }

    public JTreeExt insert(Object newChildUserObject, Object parentUserObject) {
        OptionalValue<DefaultMutableTreeNode> parent = findFirstNodeByUserObject(parentUserObject);
        if(parent.isDefined()) {
            insert(newChildUserObject, parent.get());
        } else {
            throw new RuntimeException("Cannot find parent: " + parentUserObject);
        }
        return this;
    }
    
    public DefaultMutableTreeNode insertAndGetChild(Object newChildUserObject, Object parentUserObject) {
        return this.insertAndGet(newChildUserObject, parentUserObject, true);
    }
    
    public DefaultMutableTreeNode insertAndGetParent(Object newChildUserObject, Object parentUserObject) {
        return this.insertAndGet(newChildUserObject, parentUserObject, false);
    }
    
    public DefaultMutableTreeNode insertAndGet(Object newChildUserObject, Object parentUserObject, boolean getChildOrParentNode) {
        OptionalValue<DefaultMutableTreeNode> opParent = findFirstNodeByUserObject(parentUserObject);
        if(opParent.isDefined()) {
            DefaultMutableTreeNode newChild = insert(newChildUserObject, opParent.get());
            return getChildOrParentNode ? newChild : opParent.get();
        }
        throw new RuntimeException("Cannot find parent: " + parentUserObject);
    }

    /**
     * Inserts into the selected node the given value. If there is no selection
     * @param userObject the object to add.
     * @param parent the parent node
     */
    public DefaultMutableTreeNode insert(Object userObject, DefaultMutableTreeNode parent) {
        DefaultTreeModel model = getTreeModel();
        DefaultMutableTreeNode child = new DefaultMutableTreeNode(userObject);
        model.insertNodeInto(child, parent, parent.getChildCount());
        return child;
    }

    /**
     * Returns the selected node of this tree
     * @return the selected node of this tree
     */
    public OptionalValue<DefaultMutableTreeNode> getSelectedNode() {
        TreePath selection = getSelectionPath();
        DefaultMutableTreeNode value = null;
        if(selection != null) {
            value = (DefaultMutableTreeNode) selection.getLastPathComponent();
        }
        return new OptionalValue<DefaultMutableTreeNode>(value);
    }

    /**
     * Expands the path to the given node
     * @param child the node to expand
     */
    public void expand(DefaultMutableTreeNode child) {
        DefaultTreeModel model = getTreeModel();
        TreePath path = new TreePath(model.getPathToRoot(child));
        DefaultMutableTreeNode last = (DefaultMutableTreeNode) path.getLastPathComponent();
        if(last.isLeaf() && last.getParent() != null) {
            path = new TreePath(model.getPathToRoot(last.getParent()));
        }
        expandPath(path);
    }

    /**
     * Updates the view because the given element's user object has changed
     * @param node the node whose user object has changed
     */
    public void update(DefaultMutableTreeNode node) {
        DefaultTreeModel model = getTreeModel();
        model.nodeStructureChanged(node);
        repaint();
    }

    /**
     * Adds a selection listener to this tree
     * @param listener the listener to add
     */
    public JTreeExt addSelectionListener(TreeSelectionListener listener) {
        LISTENERS.add(listener);
        getSelectionModel().addTreeSelectionListener(listener);
        return this;
    }

    /**
     * Move a node from the current parent to a new one
     * @param child the node to move
     * @param newParent the new parent of the node
     * @return this
     */
    public JTreeExt move(DefaultMutableTreeNode child, DefaultMutableTreeNode newParent) {
        DefaultTreeModel model = getTreeModel();
        model.removeNodeFromParent(child);
        model.insertNodeInto(child, newParent, newParent.getChildCount());
        return this;
    }

    /**
     * Overridden to automatically expand the root node
     */
    @Override
    public void addNotify() {
        DefaultMutableTreeNode root = getRoot();
        if(root != null) {
            expand(root);
        }
        super.addNotify();
    }

    /**
     * Applies a simple renderer the user objects of the node of this tree
     * @param labelRenderer a rendere that modifies a preset label given the rendered value
     * @return this
     */
    public JTreeExt setSimpleRenderer(final Function2<Object, JLabel, ?> labelRenderer) {
        setCellRenderer(new DefaultTreeCellRenderer() {

            private static final long serialVersionUID = 0;

            @Override
            public Component getTreeCellRendererComponent(JTree tree, Object value, boolean sel, boolean expanded, boolean leaf, int row, boolean hasFocus) {
                Component component =  super.getTreeCellRendererComponent(tree, value, sel, expanded, leaf, row, hasFocus);
                labelRenderer.apply(((DefaultMutableTreeNode) value).getUserObject(), (JLabel) component);
                return component;
            }
        });
        return this;
    }

    /**
     * Applies a simple renderer to the user object of the nodes in this tree that
     * are compatible with a given type
     * @param <T> the type of the user objects to render
     * @param userObjectType the class of the user objects to render
     * @param labelRenderer the function applied to render the objects
     * @return this
     */
    public <T> JTreeExt setSimpleRenderer(final Class<T> userObjectType, final Function2<T, JLabel, ?> labelRenderer) {
        setCellRenderer(new DefaultTreeCellRenderer() {

            private static final long serialVersionUID = 0;

            @Override
            public Component getTreeCellRendererComponent(JTree tree, Object value, boolean sel, boolean expanded, boolean leaf, int row, boolean hasFocus) {
                Component component =  super.getTreeCellRendererComponent(tree, value, sel, expanded, leaf, row, hasFocus);
                DefaultMutableTreeNode node = (DefaultMutableTreeNode) value;
                Object userObject = node.getUserObject();
                if(userObjectType.isAssignableFrom(userObject.getClass())) {
                    labelRenderer.apply(userObjectType.cast(userObject), (JLabel) component);
                }
                return component;
            }
        });
        return this;
    }

    /**
     * Returns the tree model used by this tree
     * @return the default tree model used by this tree
     */
    public DefaultTreeModel getTreeModel() {
        return (DefaultTreeModel) getModel();
    }

    /**
     * Finds the node with the given user object
     * @param matcher a function that returns true if the user object matches the find criteria
     * @return the list of nodes whose user object passes the given test function
     */
    public List<DefaultMutableTreeNode> findNodeByUserObject(Function1<Object, Boolean> matcher) {
        LinkedList<DefaultMutableTreeNode> result = new LinkedList<DefaultMutableTreeNode>();
        Enumeration<?> scan = getRoot().breadthFirstEnumeration();
        while(scan.hasMoreElements()) {
            DefaultMutableTreeNode nextElement = (DefaultMutableTreeNode) scan.nextElement();
            Object object = nextElement.getUserObject();
            if(matcher.apply(object)) result.add(nextElement);
        }
        return result;
    }

    /**
     * Iterator over the nodes of this tree
     * @return the iterator
     */
    public Iterator<DefaultMutableTreeNode> iterator() {
        final Enumeration<?> enumeration = getRoot().breadthFirstEnumeration();
        return new Iterator<DefaultMutableTreeNode>() {

            public boolean hasNext() {
                return enumeration.hasMoreElements();
            }

            public DefaultMutableTreeNode next() {
                return (DefaultMutableTreeNode) enumeration.nextElement();
            }

            public void remove() {
                throw new UnsupportedOperationException("Not supported yet.");
            }
        };
    }

    /**
     * Finds and returns the first node in the tree that contains the given user object
     * @param item the user object to find (equals)
     * @return the first node containing the given user object
     */
    public OptionalValue<DefaultMutableTreeNode> findFirstNodeByUserObject(Object item) {
        DefaultMutableTreeNode result = null;
        Enumeration<?> scan = getRoot().breadthFirstEnumeration();
        while(scan.hasMoreElements() && result == null) {
            DefaultMutableTreeNode nextElement = (DefaultMutableTreeNode) scan.nextElement();
            Object object = nextElement.getUserObject();
            if(item == null) {
                if(object == null) result = nextElement;
            } else if(item.equals(object)) {
                result = nextElement;
            }
        }
        return new OptionalValue<DefaultMutableTreeNode>(result);
    }
    
    /** Executes a function on the each elemenet of the tree until fun returns false*/
    public void foreach(Function1<DefaultMutableTreeNode, Boolean> fun) {
        Enumeration<?> scan = getRoot().breadthFirstEnumeration();
        while(scan.hasMoreElements()) {
            DefaultMutableTreeNode node = (DefaultMutableTreeNode) scan.nextElement();
            Boolean goOn = fun.apply(node);
            if(Boolean.FALSE.equals(goOn)) return;
        }
    }

    /**
     * Get the node at the given location
     * @param b the location of the node to get (relative to this jtree component)
     * @return a value holding the resulting node. Maybe undefined.
     */
    public OptionalValue<DefaultMutableTreeNode> getNodeAt(Point b) {
        DefaultMutableTreeNode node = null;
        TreePath pathForLocation = this.getPathForLocation(b.x, b.y);
        if(pathForLocation != null) {
            node = (DefaultMutableTreeNode) pathForLocation.getLastPathComponent();
        }
        return new OptionalValue<DefaultMutableTreeNode>(node);
    }

    /**
     * Set the selected node (or clears the selection if node is null9
     * @param node the node to select. Can be null. If null, clears the selection
     * @param fireSelectionEvent false to avoid the notification of a tree selection change.
     */
    public void setSelectedNode(DefaultMutableTreeNode node, boolean fireSelectionEvent) {
        DefaultTreeSelectionModel defaultTreeSelectionModel = (DefaultTreeSelectionModel) getSelectionModel();
        if(!fireSelectionEvent) {
            for (TreeSelectionListener treeSelectionListener : LISTENERS) {
                defaultTreeSelectionModel.removeTreeSelectionListener(treeSelectionListener);
            }
        }
        if(node == null) {
            getSelectionModel().clearSelection();
        } else {
            TreePath path = new TreePath(node.getPath());
            getSelectionModel().setSelectionPath(path);
        }
        if(!fireSelectionEvent) {
            for (TreeSelectionListener treeSelectionListener : LISTENERS) {
                defaultTreeSelectionModel.addTreeSelectionListener(treeSelectionListener);
            }
        }
    }

    /**
     * Remove the given node from the tree
     * @param node the node to remove
     */
    public void remove(DefaultMutableTreeNode node) {
        getTreeModel().removeNodeFromParent(node);
    }

    /**
     * Returns the selected node location or null if no node is selected
     * @return the selected node location or null
     */
    public Point getSelectedNodeLocation() {
        TreePath selection = getSelectionPath();
        if(selection == null) return null;
        Rectangle pathBounds = getPathBounds(selection);
        return pathBounds.getLocation();
    }
    
    public Point getNodeLocation(DefaultMutableTreeNode node) {
        TreeNode[] path = node.getPath();
        Rectangle pathBounds = getPathBounds(new TreePath(path));
        Point location = pathBounds.getLocation();
        location.y += pathBounds.height;
        return location;
    }

    public void selectNodeAt(Point location) {
        OptionalValue<DefaultMutableTreeNode> nodeOp = getNodeAt(location);
        setSelectedNode(nodeOp.get(), true);
    }
}
